package leetcode_1501_1600;

import node.TreeNode;

import java.util.TreeMap;

public class LeeCode_1530 {
    public static void main(String[] args) {

    }
    private static int countPairs(TreeNode root, int distance) {
        int[] ans = new int[1];
        dfs(root, distance, ans);
        return ans[0];
    }
    private static TreeMap<Integer, Integer> dfs(TreeNode node, int distance, int[] ans){
        if (node == null)
            return new TreeMap<>();
        if (node.left == null && node.right == null){
            TreeMap<Integer, Integer> map = new TreeMap<>();
            map.put(1, 1);
            return map;
        }
        TreeMap<Integer, Integer> ml = dfs(node.left, distance, ans);
        TreeMap<Integer, Integer> mr = dfs(node.right, distance, ans);
        for (Integer d1 : ml.keySet()) {
            for (Integer d2 : mr.keySet()) {
                if (d1 + d2 > distance){
                    break;
                }
                ans[0] += ml.get(d1) * mr.get(d2);
            }
        }
        TreeMap<Integer, Integer> res = new TreeMap<>();
        for (Integer key : ml.keySet()) {
            int nk = key + 1;
            res.put(nk, res.getOrDefault(nk, 0) + ml.get(key));
        }
        for (Integer key : mr.keySet()) {
            int nk = key + 1;
            res.put(key + 1, res.getOrDefault(nk, 0) + mr.get(key));
        }
        return res;
    }
}
